import time

import numpy as np
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
from sklearn import ensemble
from sklearn.metrics import mean_squared_error
import pandas as pd
import unittest
from sklearn.model_selection import StratifiedKFold  #分层抽样只用于分类任务，不适用于回归任务
from sklearn.model_selection import KFold
from joblib import dump, load

# 参考 https://blog.csdn.net/qq_32863339/article/details/100541447, https://zhuanlan.zhihu.com/p/52515873

#读取样本，并划分为X,y
def get_samples(in_path):
    df = pd.read_csv(in_path)
    X = np.array(df[df.columns[0:-1]])
    Y = np.array(df[df.columns[-1:]]).flatten()
    return X,Y

#对gbdt使用K折交叉验证
def gbdt_with_KFold(X,Y,k,n_estimators=300):
    '''
    :param X: 样本序列
    :param Y: 标签序列
    :param k: k折交叉验证
    :param n_estimators : gbdt分类器个数
    :return:
    '''
    seed = 7
    np.random.seed(seed)

    mean_mse = 0
    n_splits = k
    # skf = StratifiedKFold(n_splits=3,shuffle=True, random_state=seed)  #分层抽样只用于分类任务，不适用于回归任务
    skf = KFold(n_splits=k,shuffle=True, random_state=seed)  #StratifiedKFold是k-fold的一个变体，它会根据数据集的分布来划分，使得 划分后的数据集的目标比例和原始数据集近似。
    index = 1
    for train_index,test_index in skf.split(X,Y):

        X_train,y_train = X[train_index], Y[train_index]   #训练数据集
        X_test,y_test = X[test_index], Y[test_index]   #测试数据集

        # 实例化估计器对象
        params = {'n_estimators': n_estimators, 'max_depth': 4, 'min_samples_split': 2,
                  'learning_rate': 0.01, 'loss': 'ls'}
        gbr = ensemble.GradientBoostingRegressor(**params)

        # 估计器拟合训练数据
        gbr.fit(X_train, y_train)

        # 训练完的估计器对测试数据进行预测
        y_pred = gbr.predict(X_test)

        # 输出特征重要性列表
        # print(gbr.feature_importances_)
        mse = mean_squared_error(y_test, y_pred)
        print(f"{index} : {mean_mse}")
        index += 1
        mean_mse += mse

    # 使用当前模型参数进行重新训练
    gbr.fit(X, Y)
    dump(gbr, 'blink_EAR_thres_gbdt.joblib')
    print(f"mean_mse = {mean_mse / n_splits}")


# 将保存好的gbdt模型对输入的头部姿态（pitch，yaw，roll）进行EAR阈值推理
def gbdt_infer(model, pitch, yaw, roll):
    '''
    :param model: gbdt模型
    :param pitch:
    :param yaw:
    :param roll:
    :return:
    '''
    # 使用当前模型参数进行重新训练
    # gbr = load(model_path)
    gbr = model
    X = np.array([[pitch,yaw,roll]])
    y_pred = gbr.predict(X)
    return y_pred

#将保存好的gbdt模型对输入的头部姿态（pitch，yaw，roll）进行EAR阈值推理，计算平均耗时
def gbdt_infer_cost(model, pitch, yaw, roll):
    '''
    :param model: gbdt模型
    :param pitch:
    :param yaw:
    :param roll:
    :return:
    '''
    count = 1000
    gbr = model
    startT = time.time()
    X = np.array([[pitch, yaw, roll]])
    for i in range(0,count):
        y_pred = gbr.predict(X)
    endT = time.time()
    cost_mean = (endT -startT)/count
    return cost_mean

class MyTest(unittest.TestCase):

    def test_get_samples(self):
        X,Y = get_samples('../data/blink_train_85.csv')
        print(X,Y)

    def test_gbdt_with_KFold(self):
        X, Y = get_samples('../data/blink_train_85.csv')
        gbdt_with_KFold(X,Y,10,n_estimators=125)

    def test_gbdt_infer(self):
        model_path = 'blink_EAR_thres_gbdt.joblib'
        pitch,yaw,roll = 30,50,60
        gbr = load(model_path)
        y_pred = gbdt_infer(gbr,pitch,yaw,roll)
        print(f"y_pred = {y_pred}")

    def test_gbdt_infer_cost(self):
        model_path = 'blink_EAR_thres_gbdt.joblib'
        pitch,yaw,roll = 30,50,60
        gbr = load(model_path)
        cost = gbdt_infer_cost(gbr,pitch,yaw,roll)
        print(f"time_cost = {cost * 1000}ms")

